# -*- coding: utf-8 -*-
# 导入操作系统相关功能模块
import os
# 导入 PyTorch 深度学习框架
import torch
# 导入 PyTorch 神经网络模块
import torch.nn as nn
# 导入 PyTorch 优化器模块
import torch.optim as optim
# 从 torchvision 库导入数据集、数据转换和预训练模型模块
from torchvision import datasets, transforms, models
# 从 torch.utils.data 模块导入数据加载器
from torch.utils.data import DataLoader
# 从 torchvision.models 模块导入 MobileNet V2 模型和对应的权重
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
# 从 torch.utils.tensorboard 模块导入 SummaryWriter，用于记录训练过程
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

# 定义数据预处理转换
# 使用 transforms.Compose 将多个数据转换操作组合在一起
transform = transforms.Compose([
    # 将输入图像的大小调整为 224x224
    transforms.Resize((224, 224)),
    # 将图像数据转换为 PyTorch 张量
    transforms.ToTensor(),
    # 对图像数据进行归一化处理，使用预训练模型常用的均值和标准差
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
# 数据集目录路径
data_dir = 'D:/fruit/dataset'
# 加载训练数据集，使用之前定义的转换操作
train_dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'), transform=transform)
# 加载测试数据集，使用之前定义的转换操作
test_dataset = datasets.ImageFolder(os.path.join(data_dir, 'test'), transform=transform)

# 创建数据加载器
# 创建训练数据加载器，设置批量大小为 16，并开启数据打乱功能
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
# 创建测试数据加载器，设置批量大小为 16，不打乱数据
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)

# 初始化模型
# 加载预训练的 MobileNet V2 模型
model = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
# 获取模型分类器最后一层的输入特征数量
num_ftrs = model.classifier[1].in_features
# 修改模型分类器的最后一层，使其输出类别数与训练数据集的类别数相同
model.classifier[1] = nn.Linear(num_ftrs, len(train_dataset.classes))

# 定义损失函数和优化器
# 使用交叉熵损失函数，适用于多分类任务
criterion = nn.CrossEntropyLoss()
# 使用随机梯度下降（SGD）优化器，设置学习率为 0.001，动量为 0.9
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 初始化 TensorBoard 的 SummaryWriter
writer = SummaryWriter()

# 训练模型
# 检查是否有可用的 GPU，如果有则使用 GPU 进行训练，否则使用 CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 将模型移动到指定设备上
model.to(device)

# 训练的总轮数
num_epochs = 10
for epoch in range(num_epochs):
    # 将模型设置为训练模式
    model.train()
    # 初始化本轮训练的累计损失
    running_loss = 0.0
    # 使用 tqdm 包装 train_loader 以显示进度条
    train_loader_tqdm = tqdm(train_loader, desc=f'Epoch {epoch + 1}/{num_epochs}', unit='batch')
    for inputs, labels in train_loader_tqdm:
        # 将输入数据和标签移动到指定设备上
        inputs, labels = inputs.to(device), labels.to(device)

        # 清零优化器的梯度
        optimizer.zero_grad()

        # 前向传播，计算模型输出
        outputs = model(inputs)
        # 计算损失
        loss = criterion(outputs, labels)
        # 反向传播，计算梯度
        loss.backward()
        # 更新模型参数
        optimizer.step()

        # 累加损失
        running_loss += loss.item()

        # 更新进度条显示当前批次的损失
        train_loader_tqdm.set_postfix(loss=loss.item())

    # 计算本轮训练的平均损失
    avg_loss = running_loss / len(train_loader)
    # 使用 SummaryWriter 记录训练损失
    writer.add_scalar('Training Loss', avg_loss, epoch)
    # 打印本轮训练的信息
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss}')

# 评估模型
# 将模型设置为评估模式
model.eval()
# 初始化正确预测的样本数和总样本数
correct = 0
total = 0
# 关闭梯度计算，提高推理速度
with torch.no_grad():
    for inputs, labels in test_loader:
        # 将输入数据和标签移动到指定设备上
        inputs, labels = inputs.to(device), labels.to(device)
        # 前向传播，计算模型输出
        outputs = model(inputs)
        # 获取预测结果，返回最大值的索引
        _, predicted = torch.max(outputs.data, 1)
        # 累加总样本数
        total += labels.size(0)
        # 累加正确预测的样本数
        correct += (predicted == labels).sum().item()

# 打印模型在测试集上的准确率
print(f'Accuracy: {100 * correct / total}%')

# 保存模型
# 保存模型的状态字典
torch.save(model.state_dict(), 'model.pth')

# 关闭 SummaryWriter
writer.close()